"""
This is just a simplified skeletal framework for a genetic algorithm type of parameter
optimization.
"""

import random
from io import from_config, from_config_string
from numpy import argsort,linspace,float64, empty
import numpy
import pylab
import cPickle
import copy

class Allyl(object):
    """
    Represents one allyl, a single element in a gene.  
    """
    def __init__(self, value_set):
        self.value_set = value_set
        self.index = 0

    # ------------- VALUE ----------------------
    def __get_val(self):
        return self._value
    def __set_val(self, value):
        # if used, then self.index will not be valid
        self._value = value
    def __del_val(self):
        raise RuntimeError("Illegal attempt to delete allyl.value")
    value = property(fget=__get_val, fset=__set_val, fdel=__del_val,\
            doc= 'The value of the parameter associated with this allyl.')

    # ------------- INDEX ----------------------
    def __get_indx(self):
        return self._index
    def __set_indx(self, index):
        while index < 0:
            index += len(self.value_set)
        while index > len(self.value_set)-1:
            index -= len(self.value_set)
        self._index = index
        self.value = self.value_set[self._index]
    def __del_indx(self):
        raise RuntimeError("Illegal attempt to delete allyl.index")
    index = property(fget=__get_indx, fset=__set_indx, fdel=__del_indx,\
            doc= 'value = value_set[index]')

    def set_randomly(self):
        self.index = random.randint(0, len(self.value_set)-1)

    def __str__(self):
        rstring = "<|"+ str(self.value) + "|>"
        return rstring

    def __eq__(self,other):
        if self.index == other.index:
            return True
        else:
            return False

    def __ne__(self,other):
        return not self.__eq__(other)

class Gene():
    """
    This is a collection of allyls that together represent a parameter value set.
        It is the genes that can mutate, and cross with other genes.
    """
    def __init__(self, parameter_filename):
        # find out what parameters we have
        self.names = from_config_string("names",parameter_filename)
        self.score = None

        # set up the allyls from the parameter file
        self.allyls = []
        for name in self.names:
            value_set = from_config(name,parameter_filename)
            value_list = linspace(value_set[0],value_set[1],value_set[2])
            self.allyls.append( Allyl(value_list)  )
            self.allyls[-1].set_randomly()

    def cross_with(self, other_gene, rate):
        """
        Mix this gene up with another one.  This operation is done IN PLACE.
        Each allyl has a chance to be a crossover point.
        """
        swap = False
        for a1,a2 in zip(self.allyls, other_gene.allyls):
            # should we change the swap state?
            if random.random() < rate:
                if swap:
                    swap = False
                else:
                    swap = True
            # if we should, let's swap the allyls
            if swap:
                swap_allyl = a1
                a1 = a2
                a2 = swap_allyl

    def mutate(self, rate, random_rate):
        """
        Perhaps mutate each allyl, (randomly) changing its index.
        """
        for allyl in self.allyls:
            # randomly re-assign the index of the allyl
            if random.random() < random_rate:
                allyl.set_randomly()

            # make minor adjustments to the allyl index
            if random.random() < rate:
                if random.randint(0,1):
                    allyl.index += 1
                else:
                    allyl.index -= 1

    def __str__(self):
        rstring = str(self.score)
        for allyl in self.allyls:
            rstring += " " + str(allyl) + " "
        return rstring

    def __eq__(self,other):
        for allyl,ollyl in zip(self.allyls,other.allyls):
            if allyl != ollyl:
                return False
        return True

    def __ne__(self,other):
        return not self.__eq__(other)

    def __getitem__(self, key):
        if not -len(self.allyls) <= key < len(self.allyls):
            raise IndexError("Gene Index (%d) outside accptable range [%d, %d]" % 
                             (key,-len(self.allyls), len(self.allyls)-1) )
        else:
            return self.allyls[key].value

    def __len__(self):
        return len(self.allyls)
            
 
    def __hash__(self):
        '''
        Return the hash of a tuple of all the allyl indexes.
        '''
        value_list = []
        for allyl in self.allyls:
            value_list.append(allyl.index)
        value_tup = tuple(value_list)
        return hash(value_tup)

class Population():
    """
    A bunch of genes, along with some routines to evolve them.
    """
    def __init__(self, num_genes, initial_population_size, parameter_filename):
        self.generation_num = 0
        self.generation = [] # to keep track of all the genes ever
        self.num_genes = num_genes
        self.cross_rate         = from_config('cross_rate', parameter_filename)
        self.mutate_rate        = from_config('mutate_rate', parameter_filename)
        self.random_mutate_rate = from_config('random_mutate_rate', parameter_filename)
        self.genes = []
        self.all_genes_ever = set()
        for i in range(initial_population_size):
            # create a plausible gene
            candidate_gene = Gene(parameter_filename)
            while not self.gene_plausible(candidate_gene) or\
                  candidate_gene in self.all_genes_ever:
                candidate_gene = Gene(parameter_filename)
            # add it to the initial population
            self.genes.append( candidate_gene )
            self.all_genes_ever.add(candidate_gene)

    def import_cPickle(self, filename):
        """
        Read in the data from a previous run and replace any data we already have.
        """
        file = open(filename)
        self.generation = cPickle.load(file)
        self.genes = self.generation[-1]
        

    def evolve(self):
        # determine the scores of the model given different parameter values
        scores = []
        for gene in self.genes:
            if gene.score == None:
                # score this gene
                scores.append( self.score_model(gene) )
                # save this score with the gene 
                gene.score = scores[-1]
            else:
                # gene was already scored, so use the previous score.
                scores.append( gene.score )

        
        # find the top three genes and make them cross with the top rest of the top 50%
        sorted_indexes = argsort(scores)

        # keep track of the genes and their scores.
        genes_array = numpy.array(self.genes)
        self.generation.append( copy.deepcopy(genes_array[sorted_indexes]) )

        # make the new generation.
        top_gene    = self.genes[sorted_indexes[0]] 
        second_gene = self.genes[sorted_indexes[1]] 
        # elitism for top two genes
        new_genes = [top_gene,second_gene]
        while len(new_genes) < self.num_genes:
            #   -- determine the first parent
            first_index = abs(int( random.normalvariate(0,len(self.genes)/2) ))
            first_index = min(first_index, len(self.genes)-1)
            
            second_index = abs(int( random.normalvariate(0,len(self.genes)/2) ))
            second_index = min(second_index, len(self.genes)-1)
            # make sure parents are different.
            while second_index == first_index:
                second_index = abs(int( random.normalvariate(0,len(self.genes)/2) ))
                second_index = min(second_index, len(self.genes)-1)

            # make copies of the parents and then cross them.
            first_parent = copy.deepcopy(self.genes[ sorted_indexes[first_index] ])
            # reset scores, because it'll have to be retallied.
            first_parent.score = None
            second_parent = copy.deepcopy(self.genes[ sorted_indexes[second_index] ])
            second_parent.score = None

            first_parent.cross_with(second_parent,self.cross_rate)
            # allow the genes to mutate
            first_parent.mutate(self.mutate_rate, self.random_mutate_rate)
            second_parent.mutate(self.mutate_rate, self.random_mutate_rate)


            # only add the new genes if they're unique
            add_first_parent  = True
            add_second_parent = True
            for ngene in new_genes:
                if first_parent == ngene or not self.gene_plausible(first_parent):
                    add_first_parent = False
                    if first_parent == ngene:
                        print "first_parent is equal to ngene."
                    if not self.gene_plausible(first_parent):
                        print "first parent not plausible."
                    print first_parent
                    print ">>>>>>>>>>>>>> not adding first parent."
                if second_parent == ngene or second_parent == first_parent or \
                   not self.gene_plausible(second_parent):
                    add_second_parent = False
                    if second_parent == ngene:
                        print "second_parent is equal to ngene."
                    if second_parent == first_parent:
                        print "second_parent is equal to first_parent."
                    if not self.gene_plausible(second_parent):
                        print "second parent not plausible."
                    print second_parent
                    print ">>>>>>>>>>>>>> not adding second parent"
            if add_first_parent:
                new_genes.append(first_parent)
            if add_second_parent:
                new_genes.append(second_parent)

        self.genes = new_genes
        self.generation_num += 1

    def __str__(self):
        rstring = ""
        num = 0
        for gene in self.generation[-1]:
            num += 1
            rstring += "Gene " + str(num) + "-" + str(gene)+ "\n"
        return rstring

    def score_model(self, gene):
        """
        This is just a hook, you should overwrite it later for your specific purposes.
        """
        return 0

    def gene_plausible(self, gene):
        """
        Expected to be overwritten, returns True if newly created gene is a plausible one.
        """
        return True

    def genetic_diversity(self, gene_list):
        '''
        Returns a list containing the number of different values for each allyl in a gene set.
        '''
        # make a list of sets, one set for each allyl
        sets = []
        for allyl in gene_list[0].allyls:
            sets.append( set() )

        for gene in gene_list:
            for i in range(len(gene.allyls)):
                sets[i].add( gene.allyls[i].value )

        # now just return how many items are in each set in a list.
        return map(len, sets)

    def allyl_histograms(self, top_percent):
        '''
        Returns a list of allyl histograms for each generation for each allyl
            only the top_percent of scoring genes are included in each generation.
            0 < top_percent < 1.0
        allyl_histograms[i] is a allyl value histogram for the ith allyl
        allyl_histograms[i] is a nxm array where n is the number of possible allyl values
                    and m is the number of generations
        '''
        num_allyls = len(self.generation[-1][0].allyls)
        r = []
        for ai in range(num_allyls):
            # m = len(self.generation[-1][0].allyls[ai].value_set)
            m = 10
            n = len(self.generation)
            r.append( empty((n, m), float64) )
            # loop over generations
            for ni in range(n):
                generation = self.generation[ni]
                # generate the value histogram
                value_list = []
                score_list = []
                for gene in generation:
                    value_list.append( gene.allyls[ai].value )
                    score_list.append( gene.score )
                # make a list of the top scoring values
                top_value_indexes = argsort( score_list )[:int(len(value_list)*top_percent)]
                top_value_list = []
                for tvi in top_value_indexes:
                    top_value_list.append( value_list[tvi] )
                # actually make the histogram
                histo, edges = pylab.histogram(top_value_list, m) # use m bins    
                r[-1][ni] = histo 
            # transpose the matrix so pylab is happy when it does contour plots.
            r[-1] = r[-1].T
        return r

    def generation_scores(self):
        '''
        Returns a list of the top and average scores for each generation.
        '''
        top_scores = []
        average_scores = []
        for generation in self.generation:
            score_list = []
            for gene in generation:
                score_list.append( gene.score )
            sorted_score_indexes = argsort( score_list )
            top_scores.append( score_list[sorted_score_indexes[0]] )
            average_scores.append( sum(score_list)/len(score_list) )
        return ( top_scores, average_scores )
            
    
